"""
python fujian2_bilstm_predict.py
"""

import os
import json
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Bidirectional, LSTM, Dense
from tensorflow.keras.preprocessing.sequence import TimeseriesGenerator


def predict_with_interaction(json_folder_path, output_folder):
    # 获取所有 JSON 文件路径
    json_files = [os.path.join(json_folder_path, file) for file in os.listdir(json_folder_path) if file.endswith('.json')]

    all_data = []
    for json_file in json_files:
        with open(json_file, 'r') as f:
            data = json.load(f)
            df = pd.DataFrame(data)
            df['ds'] = pd.to_datetime(df['date'])
            df.rename(columns={'sales': 'y'}, inplace=True)
            all_data.append(df)

    # 准备数据用于训练（这里假设所有类别数据合并进行训练，以实现交互）
    all_values = np.concatenate([df['y'].values for df in all_data])
    sequence_length = 10
    generator = TimeseriesGenerator(all_values, all_values, length=sequence_length, batch_size=1)

    # 创建 BiLSTM 模型
    model = Sequential()
    model.add(Bidirectional(LSTM(50), input_shape=(sequence_length, 1)))
    model.add(Dense(1))

    # 编译模型
    model.compile(loss='mse', optimizer='adam')

    # 拟合模型
    model.fit(generator, epochs=50, verbose=0)

    for idx, df in enumerate(all_data):
        # 进行预测
        last_sequence = df['y'].values[-sequence_length:]
        predictions = []
        for _ in range((pd.to_datetime('2023-09-30') - pd.to_datetime(df['ds'].iloc[-1])).days):
            new_prediction = model.predict(last_sequence.reshape(1, sequence_length, 1))[0][0]
            predictions.append(new_prediction)
            last_sequence = np.append(last_sequence[1:], new_prediction)

        # 创建包含预测日期的数据框
        future_dates = pd.date_range(start=df['ds'].iloc[-1] + pd.Timedelta(days=1), end='2023-09-30')
        result_df = pd.DataFrame({'ds': future_dates, 'yhat': predictions})

        # 将结果转换回 JSON 格式并保存
        result_json = result_df[['ds', 'yhat']].to_json(orient='records')
        result_json = json.loads(result_json)
        output_json_path = os.path.join(output_folder, f'json_category{idx + 1}.json')
        os.makedirs(os.path.dirname(output_json_path), exist_ok=True)
        with open(output_json_path, 'w') as f:
            json.dump(result_json, f)

        # 绘制散点图并保存
        plt.figure(figsize=(10, 6))
        plt.scatter(df['ds'], df['y'], label='Original Data', c='blue')
        plt.scatter(result_df['ds'], result_df['yhat'], label='Predicted Data', c='red')
        plt.legend()
        plt.xlabel('Date')
        plt.ylabel('Sales')
        plt.title(f'BiLSTM Prediction on Category {idx + 1}')
        output_image_path = os.path.join(output_folder, f'gragh_category{idx + 1}.png')
        os.makedirs(os.path.dirname(output_image_path), exist_ok=True)
        plt.savefig(output_image_path)


if __name__ == "__main__":
    json_folder_path = r'..\fujian\fujian2\groupByCategory'
    output_folder = r'./test'
    predict_with_interaction(json_folder_path, output_folder)
